from typing import Dict, List, Optional
from instruction_re.core.datatypes import Instance, Preffix, Relation
from ..instances.MainTaskFormatter import MainTaskFormatter
from ..instances.SentenceFormatter import SenteceFormatter
from ..format_utils import add_instruction_prefix
from ..instances.Answer import AnswerFormatter
from ..instances.EntityFormatter import EntityFormatter
from ..format_utils import format_str_list_to_a_string


class REStrictTaskFormatter(MainTaskFormatter):
    """
    Task: Given sentence and entities, type the relation between them
    """

    @classmethod
    def format_data_sample(
        self,
        data_item,
        instruction: str,
        options: List[str] = None,
    ) -> Dict:

        sentence = data_item["sentence"]

        instruction = MainTaskFormatter().format_instruction_for_re_task(
            instruction, options
        )
        context = SenteceFormatter().format_sentence_for_re(data_item)

        entities_part = EntityFormatter().from_relations(data_item["relations"])
        labels = AnswerFormatter().from_relations_to_subject_verb_objext_structure(
            data_item["relations"]
        )

        result = entities_part + labels

        return {
            "instruction": instruction,
            "context": context + " " + Preffix.ANSWER.value,
            "labels": format_str_list_to_a_string(result, "<sep>"),
        }
